from matplotlib import pyplot as plt

def read_file(dataset, model, noise_type, noise_rate):
    epoch = []
    epoch_begin = 10
    sample_num = 10
    single_gradient = []
    for i in range(sample_num):
        single_gradient.append([])
    
    path = '../results/'+dataset+'/single_gradient/'+model+'/'+dataset+'_'+model+'_'+noise_type+'_'+noise_rate+'.txt'
    with open(path, 'r') as f:
        '''
        0 1 1 1 1 0 1 1 1 1 1 1 1 0 0 0 0 1 1 0 0 0 1 1
        0.11089060455560684 0.03105267882347107 0.1950453668832779 0.09061852842569351 0.23778113722801208
        '''
        lines = f.readlines()
        clean_or_not = lines[0].strip().split()
        for i, line in enumerate(lines[epoch_begin:]):
            epoch.append(i + epoch_begin)
            tmp = line.strip().split()
            for j, t in enumerate(tmp[:sample_num]):
                single_gradient[j].append(float(t))
            
    title = dataset+'_'+model+'_'+noise_type+'_'+noise_rate+'_'+str(sample_num)+'_single_gradient'
    plt.figure(figsize=(10, 10))
    ax_list = [plt.subplot(211), plt.subplot(212)]
    color = ['y', 'r', 'g', 'b', 'c', 'm', 'k', 'w']
    label = ['noise_sim', 'clean_sim']

    ax_list[0].set_xlabel('epoch')
    ax_list[0].set_ylabel('cosine sim')
    ax_list[1].set_xlabel('epoch')
    ax_list[1].set_ylabel('cosine sim')

    for i, s in enumerate(single_gradient):
        ax_list[int(clean_or_not[i])].plot(epoch, s, linewidth=1, label=label[int(clean_or_not[i])])
    plt.savefig('single_gradient/'+title+'.png')
    plt.close()

    return


read_file('cifar100', 'resnet34', 'symmetric', '0.2')
